require(ggplot2)
require(reshape2)


delta <- rho <- seq(0.1, 1, length.out = length.out)

sim_res <- list()
for (i in 1:length(delta)) {
  fn <- paste("./change/corr07/coef/res/sim_coef_change_mse_corr07_", i, ".rds", sep = '')
  sim_res[[i]] <-  readRDS(fn)
}

## data prep for plot
mse_lasso_est  <- do.call("cbind", lapply(sim_res, function(x) x[,1]))
mse_lasso_ora  <- do.call("cbind", lapply(sim_res, function(x) x[,2]))
mse_bridge     <- do.call("cbind", lapply(sim_res, function(x) x[,3]))
# mse_ridge  <- do.call("cbind", lapply(sim_res, function(x) x[,3]))
# mse_en     <- do.call("cbind", lapply(sim_res, function(x) x[,4]))

mse_lasso_dat1  <- data.frame(mse_lasso_est)
mse_lasso_dat2  <- data.frame(mse_lasso_ora)
mse_bridge_dat  <- data.frame(mse_bridge)
colnames(mse_lasso_dat1) <- colnames(mse_lasso_dat2) <- colnames(mse_bridge_dat)   <- paste("delta", delta, sep = '')

# mse_ridge_dat   <- data.frame(mse_ridge)
# mse_en_dat      <- data.frame(mse_en)

mse_lasso_dat1$rho  <- rho
mse_lasso_dat2$rho  <- rho
mse_bridge_dat$rho  <- rho

mse_lasso_dat1$method <- "HMM Lasso"
mse_lasso_dat2$method <- "Oracle Lasso"
mse_bridge_dat$method <- "HMBB"
# mse_ridge_dat$method  <- "Ridge"
# mse_en_dat$method     <- "ElasticNet"

mse_full <- rbind(mse_lasso_dat1, mse_lasso_dat2, mse_bridge_dat) #, mse_ridge_dat, mse_en_dat)

mse_full <- melt(mse_full, id.vars = c("rho", "method"), variable.name = "delta", value.name = 'L2')
mse_full$delta <- delta[as.numeric(mse_full$delta)]

mse_full$L21 <- round(mse_full$L2 / 1e1, 4)

## theoretical boundary
# boundary <- function(rho) { exp(1 - 1/rho - log(rho)) }
# delta_hat <- boundary(rho)
# delta_line <- subset(data.frame(rho, delta_hat), delta_hat > 0.05)

## plot
fig_save <- ggplot(data = mse_full, aes(x = delta, y = rho)) + geom_tile(aes(fill = L21)) +
    # scale_fill_distiller(palette = "YlGnBu", direction = 1) +
    scale_fill_distiller(palette = "RdYlBu") +
    # scale_fill_distiller(palette = "Spectral") +
    labs(x = expression(paste(delta, " = ", n / p)), y  = expression(paste(rho, " = ", k / n)), fill = expression(paste("x", 10))) +
    # geom_line(data = delta_line, aes(x = delta_hat, y = rho), color = 'gray30') +
    scale_x_continuous(expand = c(0,0)) +
    scale_y_continuous(expand = c(0,0)) +
    facet_wrap(~method) +
    ggtitle("Panel B (Normalized Estimation Loss)") +
    theme_bw() +
    theme(panel.background = element_blank(),
            panel.grid = element_blank(),
            panel.border = element_rect(color = 'gray30', size = 0.7),
            plot.title = element_text(face = 'bold'),
            strip.background = element_blank(),
            strip.text = element_text(face = 'bold', size = 12),
            axis.text.x = element_text(size = 11, color = 'gray30'),
            axis.text.y = element_text(size = 11, color = 'gray30'),
            axis.title.x = element_text(size = 12.5),
            axis.title.y = element_text(size = 12.5))

ggsave(fig_save, file = 'SI/Figure7_Panel_B.pdf', height = 2.7, width = 6)
